use crate::error::{Error, Result};
use crate::protocol::varint::{read_varint, write_varint};

#[derive(Debug, Clone)]
pub struct UdpMessage {
    pub session_id: u32,
    pub packet_id: u16,
    pub fragment_id: u8,
    pub fragment_count: u8,
    pub address: String,
    pub payload: Vec<u8>,
}

impl UdpMessage {
    pub fn new(session_id: u32, address: String, payload: Vec<u8>) -> Self {
        Self {
            session_id,
            packet_id: 0,
            fragment_id: 0,
            fragment_count: 1,
            address,
            payload,
        }
    }

    pub fn is_fragmented(&self) -> bool {
        self.fragment_count > 1
    }

    pub fn encode(&self) -> Vec<u8> {
        let mut buf = Vec::new();

        buf.extend_from_slice(&self.session_id.to_be_bytes());
        buf.extend_from_slice(&self.packet_id.to_be_bytes());
        buf.push(self.fragment_id);
        buf.push(self.fragment_count);

        write_varint(&mut buf, self.address.len() as u64);
        buf.extend_from_slice(self.address.as_bytes());

        buf.extend_from_slice(&self.payload);

        buf
    }

    pub fn decode(data: &[u8]) -> Result<Self> {
        if data.len() < 8 {
            return Err(Error::InvalidPacket);
        }

        let mut cursor = 0;

        let session_id = u32::from_be_bytes([
            data[cursor],
            data[cursor + 1],
            data[cursor + 2],
            data[cursor + 3],
        ]);
        cursor += 4;

        let packet_id = u16::from_be_bytes([data[cursor], data[cursor + 1]]);
        cursor += 2;

        let fragment_id = data[cursor];
        cursor += 1;

        let fragment_count = data[cursor];
        cursor += 1;

        let (addr_len, n) = read_varint(&data[cursor..])?;
        cursor += n;

        let address = String::from_utf8(data[cursor..cursor + addr_len as usize].to_vec())?;
        cursor += addr_len as usize;

        let payload = data[cursor..].to_vec();

        Ok(Self {
            session_id,
            packet_id,
            fragment_id,
            fragment_count,
            address,
            payload,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_udp_message_roundtrip() {
        let msg = UdpMessage::new(12345, "8.8.8.8:53".to_string(), vec![1, 2, 3, 4]);
        let encoded = msg.encode();
        let decoded = UdpMessage::decode(&encoded).unwrap();
        
        assert_eq!(decoded.session_id, 12345);
        assert_eq!(decoded.address, "8.8.8.8:53");
        assert_eq!(decoded.payload, vec![1, 2, 3, 4]);
    }
}
